from pylab import *
from scipy.io import wavfile
import matplotlib.pyplot as plt
import numpy as np
import os

if __name__ == '__main__':
    file_path = './data/'
    # file_names = os.listdir(file_path)
    # for file_name in file_names:
    samp_freq, snd = wavfile.read(file_path + 'man_howling.wav')
    snd = snd / (2. ** 15)
    single_snd = snd
    time_array = arange(0, snd.shape[0], 1)
    time_array = time_array / samp_freq
    time_array = time_array * 1000
    plt.plot(time_array, single_snd, color='k')
    plt.xlabel('time(ms)')
    plt.ylabel('amplitude')
    plt.show()

    len_snd = len(single_snd)
    fft_snd = fft(single_snd)  # 傅里叶变换
    n_unique_pts = ceil((len_snd+1)/2.0)
    fft_snd = fft_snd[0:int(n_unique_pts)]
    fft_snd = abs(fft_snd)
    fft_snd = fft_snd / float(len_snd)
    fft_snd = fft_snd ** 2

    if len_snd % 2 > 0:
        fft_snd[0:len(fft_snd)] = fft_snd[0:len(fft_snd)] * 2
    else:
        fft_snd[0:len(fft_snd)-1] = fft_snd[0:len(fft_snd)-1] * 2
    freq_array = arange(0, n_unique_pts, 1.0) * (samp_freq / len_snd)
    plt.plot(freq_array/1000, 10*log10(fft_snd), color='k')
    plt.xlabel('freqency (khz)')
    plt.ylabel('power (db)')
    plt.show()

    papr_list = []
    frame_list = []
    freq_points = []
    for i in range(int(len_snd/1024)):
        if (i+1)*1024 < len_snd:
            frame = single_snd[i*1024:(i+1)*1024]
        else:
            frame = single_snd[i * 1024:]
        pad = np.zeros((frame.shape[0],))
        frame = np.vstack([frame, pad])
        frame = np.reshape(frame, (-1,))
        frame_fft = fft(frame)
        frame_fft = frame_fft[0:int((len(frame)+1)/2.0)]
        frame_fft = abs(frame_fft)
        frame_fft = frame_fft / float(len(frame))
        frame_fft = frame_fft ** 2
        max_m_p = np.sort(frame_fft)[-5:]
        aver_p = (sum(frame_fft) - sum(max_m_p)) / len(frame_fft)
        max_p = max(frame_fft)
        freq_list = arange(0, len(frame_fft), 1.0) * (samp_freq / len(frame))
        freq_point = freq_list[np.argmax(frame_fft)] / 1000
        papr_list.append(10*log10(max_p/aver_p))
        frame_list.append(i)
        freq_points.append(freq_point)
    plt.plot(frame_list, papr_list, color='k')
    plt.xlabel('frame')
    plt.ylabel('papr')
    plt.show()
    plt.plot(frame_list, freq_points, color='k')
    plt.xlabel('frame')
    plt.ylabel('khz')
    plt.show()
